import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from model import ResNet50
from dataset import CustomDataset

"""
训练resnet模型用于训练

"""


def train(model, dataloader, criterion, optimizer):
    model.train()
    train_loss = 0
    train_acc = 0
    for inputs, labels in dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * inputs.size(0)
        preds = torch.argmax(outputs, 1)
        train_acc += torch.sum(preds == labels.data)
    train_loss = train_loss / len(dataloader.dataset)
    train_acc = train_acc.double() / len(dataloader.dataset)
    return train_loss, train_acc


def evaluate(model, dataloader, criterion):
    model.eval()
    val_loss = 0
    val_acc = 0
    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            val_loss += loss.item() * inputs.size(0)
            preds = torch.argmax(outputs, 1)
            val_acc += torch.sum(preds == labels.data)
    val_loss = val_loss / len(dataloader.dataset)
    val_acc = val_acc.double() / len(dataloader.dataset)
    return val_loss, val_acc


if __name__ == '__main__':

    """
        查看cuda是否可用
        torch.cuda.is_available()
    """
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('Using device:', device)

    """
        batch_size:每批训练的大小，将显卡跑满即可
        learning_rate:控制模型参数在每次迭代中更新的步长或速率
        num_epochs:训练代数
    """
    # Define hyperparameters
    batch_size = 32
    learning_rate = 0.001
    num_epochs = 100

    """
        (448, 448)：训练图片尺寸
        transforms.Normalize: [0.485, 0.456, 0.406] 表示图像的每个通道（RGB）的均值
                [0.229, 0.224, 0.225] 表示图像的每个通道的标准差。
    """
    # Define transforms
    transform = transforms.Compose([
        transforms.Resize((448, 448)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    """
        train_path：训练集地址
        test_path：测试集地址
        文件夹路径例子：-dataset
                        --class1
                            ---img1
                            ---img2
                        --class2
                            ---img1
                            ---img2
                            
        shuffle=True：将图片打乱，不以加载的顺序训练
        
        classes:文件夹名，也是类别名，顺序与预测结果相同
    """
    train_path = "ship/"
    test_path = "ship/"
    # Define datasets and dataloaders
    classes = ['GC', 'JZX_Z', 'none', 'SH_Q', 'SH_Z', 'YL_Q', 'YL_Z', 'ZF']
    train_dataset = CustomDataset(train_path, classes, transform=transform)
    test_dataset = CustomDataset(test_path, classes, transform=transform)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    """
    ResNetmode:选择resnet模型的大小可供选择50，101，152.数字越大模型越大，参数越多
    num_classes ： 训练的类别数，与数据集保持一致
    pre ： 是否加载预训练权重，在首次训练时打开，增量训练时关闭
    """
    # Define model, criterion and optimizer
    model = ResNet50(ResNetmode=152, num_classes=8, pre=False).to(device)

    """
    class_weights:类别在计算损失时候的权重，权重越大，类别占比越重
    """
    # 类别权重
    class_weights = [3, 2.0, 2.0, 0.5, 0.2, 0.2, 0.2, 0.2]

    # 定义损失函数，并传递类别权重
    criterion = nn.CrossEntropyLoss(weight=torch.tensor(class_weights))

    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    """
    pretrain : 预训练权重地址，用于增量训练等
    """
    pretrain = 'pth/Epoch-97-0.9905.pth'
    model.load_state_dict(torch.load(pretrain, map_location=device))

    # Train the model
    best_acc = 0
    for epoch in range(num_epochs):
        train_loss, train_acc = train(model, train_dataloader, criterion, optimizer)
        val_loss, val_acc = evaluate(model, test_dataloader, criterion)
        print(
            f'Epoch {epoch + 1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')
        if val_acc > best_acc:
            best_acc = val_acc
            modelname = f'./pth/Epoch-{str(epoch)}-{train_acc:.4f}.pth'
            torch.save(model.state_dict(), modelname)
        elif epoch == num_epochs - 1:
            torch.save(model.state_dict(), './pth/last_model.pth')
